#--------------------------------------------------------------------------
# Tool Name:  Attribute Clustering
# Source Name: AttriClustering.py
# Version: 1.0
# Author: Hu Wang(wanghugigi@gmail.com)
#
#--------------------------------------------------------------------------
#ArcGIS 9.2 + PythonWin 2.6.4 + numpy + scipy

#differences between ArcGIS 9.2 and ArcGIS 9.3
#1. polygon.centroid as a point in ArcGIS 9.3 while a unicode string in ArcGIS 9.2

#--------------------------------------------------------------------------
#Imports
#--------------------------------------------------------------------------
import HelperFunctions as HF
import numpy as np
from numpy import array, random
from scipy.cluster.vq import vq, kmeans, kmeans2, whiten, py_vq2
from scipy import linalg
from datetime import datetime
from time import strftime

def getCurTime():
    """
    get current time
    Return value of the date string format(%Y-%m-%d %H:%M:%S)
    """
    format='%Y-%m-%d %H:%M:%S'
    sdate = None
    cdate = datetime.now()
    try:
        sdate = cdate.strftime(format)
    except:
        raise ValueError
    return sdate

def get_inputs():
    """Given the list of inputs for the script when run, return an object
    containing the variables for the script"""
    try:
        '''
        sInputFC = HF.pGP.GetParameterAsText(0)                     # Input Feature Class
        sZField = HF.pGP.GetParameterAsText(1).upper()              # Attribute Field Name
        sOutputFC = HF.pGP.GetParameterAsText(2)                    # Output Feature Class
        sXY = HF.pGP.GetParameterAsText(3).upper()                  # add the XY cordinate or not
        iNumCluster = HF.pGP.GetParameterAsText(4)                  # Number of Clusters
        lsField = sZField.split(";")
        iIter = HF.pGP.GetParameterAsText(6)                        # Number of iteration
        sMethod = HF.pGP.GetParameterAsText(5)                      # method of Kmeans //Kmeans or KmeansPlusPlus
        '''
        '''
        #This block can be used to test the code using provided datasets
        sInputFC = "F:\Clustering by Python\countryClusters01.shp"
        sZField = "CASES;Population"
        sOutputFC = "F:\Clustering by Python\cobc.shp"
        sXY = "No"
        iNumCluster = "10"
        lsField = sZField.split(";")
        iIter = 10
        sMethod = "Kmeans"
        
        '''
        #This block can be used to test the code using provided datasets
        sInputFC = "G:\Hu\county_123.shp"
        sZField =  "ElevMean;Temp;Rain;Yield"
        sWeight = " 0.1875;0.25;0.25;0.3125"
        sOutputFC = "G:\Hu\county_0602_2.shp"
        sXY = "No" # Yes or No
        iNumCluster = "6"
        lsField = sZField.split(";")
        lsWeight = sWeight.split(";")        
        iIter = 1500
        sMethod = "Kmeans"
        
        #print len(lsField), len(lsWeight)

        if sXY == "Yes":
            if len(lsField) + 1 != len(lsWeight):
                raise HF.ReportError ("The format of weight does not meet specifications")
        elif sXY == "No":
            if len(lsField) != len(lsWeight):
                raise HF.ReportError ("The format of weight does not meet specifications")
        
    except: raise HF.ReportError (HF.msgParseErr)
    
    #Make a generic object, populate it with our variables, then return it
    try:
        obj = HF.Gen()
        obj.add('sInputFC',sInputFC)
        obj.add('sZField',lsField)
        obj.add('sOutputFC', sOutputFC)
        obj.add('sXY', sXY)
        obj.add('iNumCluster', iNumCluster)
        obj.add('iIter', iIter)
        obj.add('sMethod', sMethod)
        obj.add('lsWeight', lsWeight)
        

    except: raise HF.ReportError (HF.msgInputsErr)

    return obj

def normalization(lsInput):
    # type of lsInput should be ndarray
    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Normalizing the attribute data...")
    iLen = lsInput.shape
    fMax = np.amax(lsInput, axis=0)
    fMin = np.amin(lsInput, axis=0)
    fRange = fMax - fMin

    i = 0       #Row
    lsInput = lsInput.astype(float)
    lsOut = lsInput
    
    while i < iLen[1]:
        j = 0       #Line
        while j < iLen[0]:
            lsOut[j,i] = (lsInput[j,i] - fMin[i])/fRange[i]
            j = j + 1
        i = i + 1
    #print lsOut
    return lsOut

def build_value_lists():
    """Reads through feature table extracting geometry and field info.
       Returns this information in dictionary structures."""

    pRows = HF.pGP.SearchCursor(inputs.sInputFC)
    # pRows = pGP.SearchCursor(inputs.sOutputFC)
    pRow = pRows.Next()
    sKey = np.array([])
    lsBadRecs = []
    iBadRecCnt = 0
    iRec = 0 

    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage (HF.sBuildingLists)
    
    
    while pRow <> None:
        iRec = iRec + 1
        try:
            dTemp = []
            dTemp.append(pRow.GetValue("FID"))
            for Field in inputs.sZField:
                dTemp.append(pRow.GetValue(Field))
            #Add x,y of centroid
            if inputs.sXY == "Yes":
                feat = pRow.shape
                cent = feat.centroid.rsplit()
                #cent = string.split(cent, " ")
                xcent = int(float(cent[0]))
                ycent = int(float(cent[1]))
                dTemp.append(xcent)
                dTemp.append(ycent)
            sKey=np.append(sKey,dTemp)
        except:
            iBadRecCnt += 1
        pRow = pRows.Next()
    
    #print iBadRecCnt     
    # Report any problems encountered reading feature input data.
    if iBadRecCnt:
        sMessage = HF.msgReadErr % (iBadRecCnt, iRec)
        HF.pGP.AddWarning (sMessage)
        HF.pGP.AddWarning(`lsBadRecs`)

    if inputs.sXY == "Yes":
        sKey.shape = (-1,len(inputs.sZField)+3)
    else:
        sKey.shape = (-1,len(inputs.sZField)+1)
    
    #print sKey    
    return sKey


def kinit (X, k):
    'init k seeds according to kmeans++'
    n = X.shape[0]

    'choose the 1st seed randomly, and store D(x)^2 in D[]'
    centers = [X[random.randint(n)]]
    D = [linalg.norm(x-centers[0])**2 for x in X]

    for _ in range(k-1):
        bestDsum = bestIdx = -1

        for i in range(n):
            'Dsum = sum_{x in X} min(D(x)^2,||x-xi||^2)'
            Dsum = reduce(lambda x,y:x+y,
                          (min(D[j], linalg.norm(X[j]-X[i])**2) for j in xrange(n)))

            if bestDsum < 0 or Dsum < bestDsum:
                bestDsum, bestIdx  = Dsum, i

        centers.append (X[bestIdx])
        D = [min(D[i], linalg.norm(X[i]-X[bestIdx])**2) for i in xrange(n)]

    return array (centers)

def choose_initial_plusplus(data, k, distfun=None):
    """Choose k initial *different* centroids randomly using the
    k-means++ euristic by David Arthur and Sergei Vassilvitskii.
    This often gives better clustering results, but it is slower than the
    basic choice of starting points."""
    # See article "k-means++: The Advantages of Careful Seeding" by
    #   David Arthur and Sergei Vassilvitskii.
    # See also: http://theory.stanford.edu/~sergei/kmeans
    from random import choice, randrange, random
    from itertools import izip
    from bisect import bisect

    def weigthed_choice(objects, frequences):
        if len(objects) == 1:
            return 0

        addedFreq = []
        lastSum = 0
        for freq in frequences:
            lastSum += freq
            addedFreq.append(lastSum)

        return bisect(addedFreq, random())

    if distfun is None:
        def distfun(p1, p2):
            return sum( (cp1 - cp2) * (cp1 - cp2) for cp1, cp2 in izip(p1, p2) )

    # choose an intial centroid randomly
    #data = []
    data = data[:]
    #print type(data)
    pos = randrange(len(data))
    centroids = set([tuple(data[pos])])
    #del data[pos] # slow
    data = np.delete(data, pos, axis = 0)

    ntries = 0
    while len(centroids) < k and ntries < (k * 5):
        min_dists = [min( distfun(c, x) for c in centroids) for x in data]
        tot_dists = float(sum(min_dists))
        probabilities = [min_dists[i] / tot_dists for i, x in enumerate(data)]
        pos = weigthed_choice(data, probabilities) # this can be made faster
        centroids.add(tuple(data[pos]))
        ntries += 1
        #del data[pos] # slow
        data = np.delete(data, pos, axis = 0)

    result = map(list, centroids)
    if len(result) < k:
        # Fill in missing centroids
        result.extend( [result[0]] * (k - len(result)) )

    #npResult = np.array([])
    #npResult = result
    return result

def output_results(Cluster):
    """ This function writes results to the output feature class. """
    
    # Add results field.
    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage ("Adding results fields...")
    addField = []
    
    for Field in inputs.sZField:
        if len(Field) > 8:
            Field = Field[0:8]
        #Field = Field + "_N"
        addField.append(Field + "_N")

    if inputs.sXY == "Yes":
        addField.append("X_cent_N")
        addField.append("Y_cent_N")
    
    for Field in addField:
        if not properties.dcFields.has_key(Field.upper()):
            HF.pGP.AddField(inputs.sOutputFC, Field, "FLOAT")

    outField = "CLUSTER"    
    if not properties.dcFields.has_key(outField.upper()):
        HF.pGP.AddField(inputs.sOutputFC, outField, "LONG")
    else:
        outField = "CLUSTER_1"
        if not properties.dcFields.has_key(outField.upper()):
            HF.pGP.AddField(inputs.sOutputFC, outField, "LONG")
        
    #addField.append("CLUSTER") 

    # Add results to output FC
    HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
    HF.pGP.AddMessage (HF.sWritingResults)
    pRows = HF.pGP.UpdateCursor(inputs.sOutputFC)
    pRow = pRows.Next()
    iError = 0
    iCnt = 0
    fInterval = len(attri) / 5.0
    fMore = fInterval
    iComplete = 20
    
    while pRow <> None:
        iKey = pRow.GetValue("FID")

        #print iKey
        try:
            if Cluster[iKey]:  # make sure we have a non-Null result.
                #since the type of Cluster[iKey] is numpy.int32, it should convert to [int]
                for Field in addField:
                    iIndex = addField.index(Field)
                    temp = float(clusterAttri_Norm[iKey, iIndex])
                    pRow.SetValue(Field, temp)

            temp = int(Cluster[iKey])
            pRow.SetValue(outField, temp)
                    
            pRows.UpdateRow(pRow)
            iCnt = iCnt + 1
            if iCnt > fInterval: fInterval, iComplete = HF.check_progress(fInterval, fMore, iComplete)
        #except: pass
        except:
            iError = iError + 1
        
        pRow = pRows.Next()
        

    #print "++++++++++++++++++++++"
    print iError
    
        
    HF.pGP.AddMessage (HF.s100Percent)
    HF.pGP.AddMessage(" ")
    pRows = None           
    #return sField  



#--------------------------------------------------------------------------
#MAIN

if __name__ == "__main__":
       
    inputs = get_inputs()
    iNumCluster = int(inputs.iNumCluster)
    iIter = int(inputs.iIter)


    # set weight
    fWeight = np.array([],dtype=float)
    for item in inputs.lsWeight:
        fWeight = np.append(fWeight, float(item))
    if inputs.sXY == "Yes":
        fWeight[-1] = fWeight[-1]/2
        fWeight = np.append(fWeight, fWeight[-1])

    if sum(fWeight) >1.02 or sum(fWeight) <0.98:
        raise HF.ReportError ("The sum of weight doesn't equal to 1")
    #print fWeight
     
    
    #if iNumRecs < 30:
    #    HF.pGP.AddWarning (msgFewRecsWrn)
    
    if HF.pGP.exists(inputs.sOutputFC):
        HF.pGP.delete(inputs.sOutputFC)

    #Copy the input feature class to the output feature class.
    try:
        HF.pGP.QualifiedFieldNames = 0
        HF.pGP.Copyfeatures(inputs.sInputFC, inputs.sOutputFC)
    except:
        sMessage = HF.msgOverwriteErr % (inputs.sOutputFC)
        raise HF.ReportError (sMessage)

    properties = HF.get_featureclass_properties(inputs.sOutputFC)    
    
    attri = build_value_lists()
     
    #HF.pGP.AddMessage ("--------build_value_lists finished--------")
    clusterAttri = attri[:,1:]

    #print clusterAttri
   
    clusterAttri_Norm = normalization(clusterAttri)

    #print "----------------------"
    #print clusterAttri_Norm
    
    #KData is the clusterAttri plus the sqrt of weight
    
    KData = clusterAttri_Norm * np.sqrt(fWeight)
    

    #print "----------------------"
    #print KData

    
    #HF.pGP.AddMessage ("--------normalization finished--------")
    if inputs.sMethod == "Kmeans":
        HF.pGP.AddMessage ("--------" + getCurTime() + "--------")
        HF.pGP.AddMessage (HF.sCalculating + "( kmeans )")
        res,idx = kmeans2(KData, iNumCluster, iter =iIter)
        HF.pGP.AddMessage ("--------Kmeans finished--------")
    elif inputs.sMethod == "KmeansPlusPlus":
        HF.pGP.AddMessage ("--------" + getCurTime() + "--------")  
        HF.pGP.AddMessage (HF.sCalculating + "( centers )")
        print getCurTime()
        print "start calculating initial centers..."
        centers = choose_initial_plusplus(KData,iNumCluster)
        centers = np.array(centers)
        HF.pGP.AddMessage ("--------" + getCurTime() + "--------")  
        HF.pGP.AddMessage (HF.sCalculating + "( kmeans++ )")
        print getCurTime()
        print "start calculating kmeans++..."
        res,idx = kmeans2(KData, centers, minit='points')
    else:
        raise HF.ReportError ("Error in choosing method!")

    
    idx = idx + 1
    

    
    '''print clusterAttri_Norm[0,:]

    print "res ="
    print res[0,:]
    res = res[0:1,:]
    
    code, mind_dist = py_vq2(clusterAttri_Norm, res)

    print sum(mind_dist)
    
    
    print "code ="
    print code
    print "mind_dist ="
    print mind_dist
    '''
    '''
    idx.shape = (-1,1)
    attri = np.append(attri, idx, axis = 1)
    idx = np.arange(900)'''
    
    output_results(idx)







    